import open3d as o3d
import numpy as np
from scipy import optimize
np.set_printoptions(precision=4, suppress=True)

import pdb
def showPCDRGB(l_p,xyz,rgb,size_frame=1):
    # vis = o3d.visualization.Visualizer()
    # vis.create_window()
    point_cloud_o3d = o3d.geometry.PointCloud()
    point_cloud_o3d.points = o3d.utility.Vector3dVector(xyz[:,:3])
    point_cloud_o3d.colors = o3d.utility.Vector3dVector(rgb/255)
    frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=size_frame)
    # vis.add_geometry(point_cloud_o3d)
    # vis.add_geometry(frame)
    pcd = o3d.io.read_point_cloud(l_p)

    vis_add = o3d.geometry.PointCloud()
    vis_add.points = o3d.utility.Vector3dVector([[0,0,1]])
    vis_add.colors = o3d.utility.Vector3dVector([[0,0,1]])

    o3d.visualization.draw_geometries([point_cloud_o3d,frame,pcd])
    # vis.run()
    # vis.destroy_window()

def r(x, y, xc, yc):
    return np.sqrt((x-xc)**2 + (y-yc)**2)

def f(c, x, y):
    Ri = r(x, y, *c)
    return np.square(Ri - Ri.mean())

def least_squares_circle(coords):
    """
    Circle fit using least-squares solver.
    Inputs:

        - coords, list or numpy array with len>2 of the form:
        [
    [x_coord, y_coord],
    ...,
    [x_coord, y_coord]
    ]
        or numpy array of shape (n, 2)

    Outputs:

        - xc : x-coordinate of solution center (float)
        - yc : y-coordinate of solution center (float)
        - R : Radius of solution (float)
        - residu : MSE of solution against training data (float)
    """

    x, y = None, None
    if isinstance(coords, np.ndarray):
        x = coords[:, 0]
        y = coords[:, 1]
    elif isinstance(coords, list):
        x = np.array([point[0] for point in coords])
        y = np.array([point[1] for point in coords])
    else:
        raise Exception("Parameter 'coords' is an unsupported type: " + str(type(coords)))

    # coordinates of the barycenter
    x_m = np.mean(x)
    y_m = np.mean(y)
    center_estimate = x_m, y_m
    center, _ = optimize.leastsq(f, center_estimate, args=(x, y))
    xc, yc = center
    Ri       = r(x, y, *center)
    R        = Ri.mean()
    residu   = np.sum((Ri - R)**2)
    return xc, yc, R, residu

def get_center_xyz(cylinder_pl,pcd):
    # cylinder_pl = [2.5,3.5]
    z_range=[0.2,2]

    x_range = np.array([cylinder_pl[0]-0.3,cylinder_pl[0]+0.3])
    y_range = np.array([cylinder_pl[1]-0.3,cylinder_pl[1]+0.3])
    

    crop_box = o3d.geometry.AxisAlignedBoundingBox([x_range[0], y_range[0], z_range[0]], [x_range[1],y_range[1], z_range[1]])
    pcd_cropped = pcd.crop(crop_box)

    # o3d.visualization.draw_geometries([pcd_cropped])

    # cylinder
    xyz = np.asarray(pcd_cropped.points)
    a, b, r, residual = least_squares_circle(xyz[:,:2])
    print('Center: [{}, {}, {}], Radius: {}'.format(a, b, 0, r))
    cylinder = o3d.geometry.TriangleMesh.create_cylinder(radius=r, height=1)
    cylinder.translate([a,b,0])

    frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=2)
    o3d.visualization.draw_geometries([pcd_cropped,cylinder,frame])
    print("xy: ", [a,b])
    return a,b

def pick_points(pcd,size_frame=0.3):
    print("")
    print(
        "1) Please pick at least three correspondences using [shift + left click]"
    )
    print("   Press [shift + right click] to undo point picking")
    print("2) After picking points, press 'Q' to close the window")
    vis = o3d.visualization.VisualizerWithEditing()
    vis.create_window(window_name="pick pts",visible=True)
    frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=size_frame)
    vis.add_geometry(pcd)
    vis.add_geometry(frame)
    # o3d.visualization.draw_geometries([point_cloud_o3d,frame])
    
    vis.run()  # user picks points
    vis.destroy_window()
    print("")
    return vis.get_picked_points()

def show_plane(a,b,c,d,crop_pcd,inliers):

    plane = o3d.geometry.TriangleMesh()

    vertices = [
        [100.0, 20.0, (-d - a - b)],
        [-100.0, 20.0, (-d + a + b)],
        [-100.0, -20.0, (-d + a - b)],
        [100.0, -20.0, (-d - a + b)]
    ]
    plane.vertices = o3d.utility.Vector3dVector(vertices)

    triangles = [
        [0, 1, 2],
        [0, 2, 3]
    ]
    plane.triangles = o3d.utility.Vector3iVector(triangles)
    plane.paint_uniform_color([0, 1, 0])
    frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=2)

    # normal
    normal = [a, b, c]
    origin = [0, 0, -d/c]
    length = 2
    end_point = [origin[i] + length * normal[i] for i in range(3)]
    line_set = o3d.geometry.LineSet()
    line_set.points = o3d.utility.Vector3dVector([origin, end_point])
    line_set.lines = o3d.utility.Vector2iVector([[0, 1]])
    line_set.colors = o3d.utility.Vector3dVector([[1, 1, 0]])
    
    
    # normal
    normal = [0, 0, 1]
    origin = [0, 0, 0]
    length = 2
    end_point = [origin[i] + length * normal[i] for i in range(3)]
    line_set_z = o3d.geometry.LineSet()
    line_set_z.points = o3d.utility.Vector3dVector([origin, end_point])
    line_set_z.lines = o3d.utility.Vector2iVector([[0, 1]])
    line_set_z.colors = o3d.utility.Vector3dVector([[1, 0, 0]])

    vis = o3d.visualization.Visualizer()
    vis.create_window()
    vis.add_geometry(plane)
    # vis.add_geometry(frame)
    vis.add_geometry(line_set)
    vis.add_geometry(line_set_z)
    # vis.add_geometry(crop_pcd)
    
    # pdb.set_trace()
    point_cloud_o3d = o3d.geometry.PointCloud()
    point_cloud_o3d.points = o3d.utility.Vector3dVector(inliers[:,:3])
    point_cloud_o3d.colors = o3d.utility.Vector3dVector(inliers/255)
    vis.add_geometry(point_cloud_o3d)

    vis.run()
    vis.destroy_window()

def get_plane(pcd,debug=True):
    bbox=o3d.geometry.AxisAlignedBoundingBox([-50, -20, -0.5], [100,20, 0.5])
    clipped_pcd = pcd.crop(bbox)

    plane_model, inliers = clipped_pcd.segment_plane(distance_threshold=0.02,
                                                    ransac_n=3,
                                                    num_iterations=1000)

    [a, b, c, d] = plane_model
    inlier_points = np.asarray(clipped_pcd.points)[inliers]

    # if debug:
    #     show_plane(a,b,c,d,clipped_pcd,inlier_points)
    print("plane: {} {} {} {} ".format(a, b, c, d))
    return a,b,c,d

def calculate_z_given_plane_equation(a, b, c, d, x, y):
    z = (-d - a*x - b*y) / c
    return z

def get_lidar_pts_manual(lidar_p,pts_num):
    pcd = o3d.io.read_point_cloud(lidar_p)

    # ground fit , get z of p-cylinder in plane
    a,b,c,d = get_plane(pcd)

    # pts_num = 4 cylinder center
    lidar_pts = []
    for i in range(pts_num):
        pts_ind = pick_points(pcd)
        circle_x,circle_y,_ = pcd.points[pts_ind[0]]
        x,y = get_center_xyz([circle_x,circle_y],pcd)

        z = calculate_z_given_plane_equation(a,b,c,d,x,y)
        print("x,y,z \n",x,y,z)
        lidar_pts.append([x,y,z])

    return np.array(lidar_pts,dtype=np.float32)

def get_lidar_pts_manual_not_fit(lidar_p,pts_num):
    pcd = o3d.io.read_point_cloud(lidar_p)

    # ground fit , get z of p-cylinder in plane
    a,b,c,d = get_plane(pcd)

    # pts_num = 4 cylinder center
    lidar_pts = []
    for i in range(pts_num):
        pts_ind = pick_points(pcd)
        circle_x,circle_y,_ = pcd.points[pts_ind[0]]
        x,y = circle_x,circle_y

        z = calculate_z_given_plane_equation(a,b,c,d,x,y)
        print("x,y,z \n",x,y,z)
        lidar_pts.append([x,y,z])

    return np.array(lidar_pts,dtype=np.float32)



def get_lidar_pts_auto(lidar_p,pts_num,xyz_min,xyz_max,debug=True):
    pcd = o3d.io.read_point_cloud(lidar_p)
    a,b,c,d = get_plane(pcd,debug)

    # ground fit , get z of p-cylinder in plane
    # get 4 cylinder cluster
    # crop
    # xyz_min=[6, -5, 0.1]
    # xyz_max=[15,5, 1]
    bbox=o3d.geometry.AxisAlignedBoundingBox(xyz_min,xyz_max)
    clipped_pcd = pcd.crop(bbox)
    if debug:
        frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=2)
        o3d.visualization.draw_geometries([clipped_pcd,frame])

    # cluster 
    with o3d.utility.VerbosityContextManager(o3d.utility.VerbosityLevel.Debug) as cm:
        labels = np.array(clipped_pcd.cluster_dbscan(eps=0.4, min_points=10, print_progress=True))

    print("the number of cylinder cluster is ",len(np.unique(labels)))
    if len(np.unique(labels))!= pts_num:
        print("Please check crop range!,the number of cylinder cluster != ",pts_num)
        exit()
        
    lidar_pts = []
    i=0
    for cluster_id in np.unique(labels):
        if cluster_id == -1:
            continue
        xyz = np.asarray(clipped_pcd.points)[labels == cluster_id]
        x, y, r, residual = least_squares_circle(xyz[:,:2])

        z = calculate_z_given_plane_equation(a,b,c,d,x,y)
        # print('Center: [{}, {}, {}], Radius: {}'.format(x,y,z, r))
        lidar_pts.append([x,y,z,r])
        i+=1
    
    lidar_pts = np.array(lidar_pts,dtype=np.float32)
    print("lidar_pts: ",lidar_pts)
    # sort by xy
    colors = [[1, 0, 0], [0, 1, 0], [0, 0, 1],[0, 0, 0]]
    '''
    *   *
    p4(black)  p3(blue)
    *   *
    p1(red)  p2(green)
    '''
    lidar_pts_id = np.array([1,2,3,4])
    lidar_pts_w_id = np.hstack([lidar_pts,lidar_pts_id.reshape(4,1)])
    
    column_means = np.mean(lidar_pts_w_id, axis=0)

    for i in range(len(lidar_pts_w_id)):
        x,y,z,r,id = lidar_pts_w_id[i]
        if x<column_means[0] and y>column_means[1]:
            lidar_pts_w_id[i][-1]=1
        if x<column_means[0] and y<column_means[1]:
            lidar_pts_w_id[i][-1]=2
        if x>column_means[0] and y<column_means[1]:
            lidar_pts_w_id[i][-1]=3
        if x>column_means[0] and y>column_means[1]:
            lidar_pts_w_id[i][-1]=4
    sorted_indices = np.argsort(lidar_pts_w_id[:, -1])
    # print("cylinder sort int: ", sorted_indices)
    lidar_pts_w_id = lidar_pts_w_id[sorted_indices]
    print("cylinder sort: \n", lidar_pts_w_id)
    print("range: ",bbox)
    if debug:
        pcd_cylinder_show = []
        for i in range(len(lidar_pts_w_id)):
            x,y,z,r,ID = lidar_pts_w_id[i]

            cylinder = o3d.geometry.TriangleMesh.create_cylinder(radius=r, height=1)
            cylinder.translate([x,y,0])
            cylinder.paint_uniform_color(colors[i])
            pcd_cylinder_show.append(cylinder)

        frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=2)
        pcd_cylinder_show.append(frame)
        pcd_cylinder_show.append(clipped_pcd)
        o3d.visualization.draw_geometries(pcd_cylinder_show)

    pts = np.array(lidar_pts_w_id[:,:3],dtype=np.float32)
    return pts

if __name__=="__main__":
    lidar_p = "/home/westwell/welldriver/log/data_calib/qtruck_psa001/zjh_0307_snap/snap_rl/qtruck_psa001_rl/lidar_left_ego/1709787433725392000.pcd"
    pts_num = 2
    pts = get_lidar_pts_manual(lidar_p,pts_num)
    print(pts)
    # [[2.4653146  3.6818726  0.        ]
    # [3.0435183  4.168749   0.        ]
    # [0.04880002 4.770845   0.        ]
    # [2.8511071  6.6054993  0.        ]]
    




